import os
from langchain.tools import BaseTool
from pydantic import BaseModel,Field,ConfigDict
from typing import Any,Type
from langchain_chroma import Chroma
from model.my_chat_model import ChatModel
#定义输入参数的数据模型类型，... 表示必填
class ChromaInput(BaseModel):
    question:str = Field(...,description="企业背景查询问题")

#定义工具类
class ChromaTool(BaseTool):
    # 定义模型，是否允许输入参数，输入参数的数据模型类
    model_config = ConfigDict(arbitrary_types_allowed=True)
    def __init__(self,**kwargs:Any):
        super().__init__(
            name = "get_chroma_tool",
            description="主要用于查询chroma数据库的企业背景信息,用户必须输入的参数是问题",
            **kwargs
        )

    #定义工具参数
    args_schema : Type[BaseModel] = ChromaInput

    def _run(self,question:str):
        # 1. 获取 embedding 模型（用于 Chroma 检索）
        chat = ChatModel()
        embedding_model = chat.get_embedding_model()
        # 服务器环境的路径
        db = os.path.join("chroma_db")
        # 本地测试地址
        # db = "../chroma_db"
        # print(f"db={db}")

        # 2. 加载 Chroma 向量库
        store = Chroma(
            persist_directory=db,
            embedding_function=embedding_model,
            collection_name="companies_context",
            collection_metadata={"hnsw:space":"cosine"}
        )
        # 3. 初步检索 top-10 文档
        # search_type="similarity" 根据文档的余玄相似度进行查询，k表示返回的文档数量
        res=store.as_retriever(search_type="similarity",search_kwargs={"k":1})
        #检索
        docs=res.invoke(question)
        print(f"docs:{docs}")
        print(f"初步检索到 {len(docs)}个文档")
        if not docs:
            return "未检索到相关企业背景信息"

        # 7. 拼接最终上下文和来源
        context_text = "\n\n".join([doc.page_content.strip() for doc in docs])

        #去重来源文件
        source_texts = list(set(doc.metadata.get("source","未知") for doc in docs))
        source_text = "\n".join(source_texts)
        whole_text = f"{context_text}\n\n数据来源:\n{source_text}"
        return whole_text

# if __name__=="__main__":
#     tool = ChromaTool()
#     text=tool._run(question="请查询一下公司2的背景信息")
#     print(text)


